package com.cmict.data.system.service.websocket.component;


import com.cmict.data.system.dal.dataobject.oauth2.OAuth2AccessTokenDO;
import com.cmict.data.system.dal.redis.oauth2.OAuth2AccessTokenRedisDAO;
import com.cmict.data.system.service.websocket.constant.AdminWebSocketConstants;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.http.server.ServletServerHttpRequest;
import org.springframework.http.server.ServletServerHttpResponse;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.Map;

/**
 * WebSocket 握手拦截器 在握手时记录下当前 session 对应的用户Id和token信息
 *
 * @author Hccake 2021/1/4
 * @version 1.0
 */
@Slf4j
@RequiredArgsConstructor
public class UserAttributeHandshakeInterceptor implements HandshakeInterceptor {
	@Resource
	private OAuth2AccessTokenRedisDAO loginUserRedisDAO;

	/**
	 * Invoked before the handshake is processed.
	 * @param request the current request
	 * @param response the current response
	 * @param wsHandler the target WebSocket handler
	 * @param attributes the attributes from the HTTP handshake to associate with the
	 * WebSocket session; the provided attributes are copied, the original map is not
	 * used.
	 * @return whether to proceed with the handshake ({@code true}) or abort
	 * ({@code false})
	 */
	@Override
	public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
			Map<String, Object> attributes) {
		String accessToken = null;
		// 获得 accessToken
		if (request instanceof ServletServerHttpRequest) {
            HttpServletRequest serverRequest = ((ServletServerHttpRequest) request).getServletRequest();
            accessToken = serverRequest.getHeader("Sec-WebSocket-Protocol");
		}
		// 由于 WebSocket 握手是由 http 升级的，携带 token 已经被 Security 拦截验证了，所以可以直接获取到用户
		OAuth2AccessTokenDO redisLoginUser = loginUserRedisDAO.get(accessToken);
		if(redisLoginUser != null){
			attributes.put(AdminWebSocketConstants.TOKEN_ATTR_NAME, accessToken);
			attributes.put(AdminWebSocketConstants.USER_KEY_ATTR_NAME,redisLoginUser.getTenantId() +"_"+redisLoginUser.getUserId());
			log.info("websocket 租户为 "+ redisLoginUser.getTenantId() +"       websocket  用户id 为 "+redisLoginUser.getUserId());
			return true;
		}else {
			return false;
		}

//		attributes.put(AdminWebSocketConstants.USER_KEY_ATTR_NAME, "1");

	}

	/**
	 * Invoked after the handshake is done. The response status and headers indicate the
	 * results of the handshake, i.e. whether it was successful or not.
	 * @param request the current request
	 * @param response the current response
	 * @param wsHandler the target WebSocket handler
	 * @param exception an exception raised during the handshake, or {@code null} if none
	 */
	@Override
	public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, WebSocketHandler wsHandler,
			Exception exception) {
        if (request instanceof ServletServerHttpRequest) {
            ServletServerHttpRequest serverRequest = (ServletServerHttpRequest) request;
            String accessToken = serverRequest.getServletRequest().getHeader("Sec-WebSocket-Protocol");
            if(StringUtils.isNotEmpty(accessToken)){
                HttpServletResponse httpResponse = ((ServletServerHttpResponse)response).getServletResponse();
                httpResponse.addHeader("Sec-WebSocket-Protocol",accessToken);
            }
        }
		// doNothing
	}

}
